#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : lei.wang
import asyncio
from typing import Annotated
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import StreamingResponse
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.db.database import async_get_db
from app.services.user import user_service
from app.schemes.prompts import UserPromptsModel
from app.core.pkg.log import logger
from app.util.common import ask_birthdate, process_question
from app.schemes.user import UserCreate
from app.services.chat_messages import ChatMessagesService
from app import config
from app.util.jwt_helper import validate_auth_header, validate_user_score
ask_router = APIRouter(prefix="/ask", tags=["ASK"])


async def stream_response(generator):
    """流式响应生成器，用于将AI响应内容流式传输到客户端"""
    try:
        # 确保生成器是可迭代的
        if hasattr(generator, '__aiter__'):
            async for text in generator:
                yield f"data: {text}\n\n"
        else:
            # 如果不是异步迭代器，尝试直接获取结果
            result = await generator
            if isinstance(result, str):
                yield f"data: {result}\n\n"
            else:
                # 处理其他可能的返回类型
                async for text in result:
                    yield f"data: {text}\n\n"
    except Exception as e:
        logger.error(f"流式响应生成出错: {str(e)}")
        yield f"data: 处理出错: {str(e)}\n\n"
    finally:
        yield "data: [DONE]\n\n"


@ask_router.post('/start', description='开始问答')
async def ask_bazi_question(request: Request, ask_model: UserPromptsModel, db: Annotated[AsyncSession, Depends(async_get_db)]):
    """
    提供问答服务：
    1. 根据用户ID查询用户出生信息，如果查询不到则调用 ask_birthdate 提示用户提供出生信息；
    2. 如果收到符合格式的出生信息，则将信息写入数据库，并基于出生信息生成命理数据后调用大模型回答问题。
    3. 支持流式返回，使用SSE格式向客户端推送数据。
    """
    user_palyload_id = validate_auth_header(request)
    await validate_user_score(user_palyload_id, db)
    logger.info(f"接收到 /ask 请求，用户ID: {user_palyload_id}，问题: {ask_model.question}")
    # 首先查询用户出生信息
    user_info = await user_service.get(db=db, payload_id=user_palyload_id)

    # 如果查不到（此时用户的话里可能已有出生年月日时，也可能没有）用户认证后一定存在用户信息
    if not user_info:
        logger.info(f"用户({user_palyload_id})出生信息未查询到，调用 ask_birthdate。")
        response = await asyncio.to_thread(ask_birthdate, ask_model.question)
        # 如果用户说的话中有出生信息，此时大模型的回复会有体现，直接提取出生信息和回答，并将数据插入数据库
        if "&&" in response:
            # 用户提供了正确格式的信息(下面处理response的格式的原因跟prompt有关)
            user_birthdate = response.split("&&")[0]
            answer = response.split("&&")[1]
            birthdate_part = user_birthdate.split("#")[0]
            sex_part = user_birthdate.split("#")[1]
            obj = UserCreate(
                birthday=birthdate_part,
                payload_id=ask_model.payload_id,
                gender=sex_part,
            )
            await user_service.create(db=db, object=obj)
            logger.info(f"用户({user_palyload_id})的出生信息已保存，返回回答。")

            # record messages
            await ChatMessagesService.create_chat(db, user_palyload_id, ask_model.question, answer)
            
            # 将完整答案转换为单个元素的流返回
            async def answer_generator():
                yield answer
            
            return StreamingResponse(
                stream_response(answer_generator()),
                media_type="text/event-stream"
            )
        else:
            # 如果用户说的话中没有出生信息，直接返回提示。
            logger.info(f"用户({user_palyload_id})未提供有效的出生信息，返回提示。")
            # record messages
            await ChatMessagesService.create_chat(db, user_palyload_id, ask_model.question, response)
            
            # 将提示信息转换为单个元素的流返回
            async def response_generator():
                yield response
            
            return StreamingResponse(
                stream_response(response_generator()),
                media_type="text/event-stream"
            )
    else:
        # 如果查询到了用户的出生信息则会直接回复
        try:
            # 从数据库中取出的出生信息格式： [id, birthdate, sex, time]
            # 处理日期字符串，确保替换T为空格，并移除时区信息+08:00
            birthdate = user_info.get('birthday', '')
            if '+' in birthdate:
                birthdate = birthdate.split('+')[0]  # 移除时区信息如 +08:00
            birthdate = birthdate.replace("T", " ")
            logger.debug(f"用户({user_palyload_id})查询到的出生信息：birthdate={birthdate}, gender={user_info.get('gender', '')}")
            
            # 使用流式处理
            async def stream_ai_response():
                ai_response_gen = await process_question(birthdate, user_info.get('gender', ''), ask_model.question, db, user_palyload_id, stream=True)
                async for chunk in ai_response_gen:
                    yield chunk
            
            logger.info(f"用户({user_palyload_id})/ask 使用流式返回处理请求。")
            
            return StreamingResponse(
                stream_response(stream_ai_response()),
                media_type="text/event-stream"
            )
        except Exception as e:
            logger.error(f"处理用户({user_palyload_id}) /ask 请求时出错：{str(e)}")
            
            # 错误信息也通过流返回
            async def error_generator():
                yield f"内部错误：{str(e)}"
                
            return StreamingResponse(
                stream_response(error_generator()),
                media_type="text/event-stream"
            )


@ask_router.get('/history', description=f'获取{config.get("chat_round").get("number")}轮历史记录')
async def get_chat_message_round(request: Request, db: Annotated[AsyncSession, Depends(async_get_db)]):
    user_palyload_id = validate_auth_header(request)
    result = await ChatMessagesService.get_chat_messages(db=db, payload_id=user_palyload_id)
    return {"success": True, "message": "success", **result}
